import websocket  # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
import uuid
import json
import requests
import os
import sys
from datetime import datetime

server_address = "127.0.0.1:8188"
client_id = str(uuid.uuid4())

def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    headers = {'Content-Type': 'application/json'}
    response = requests.post(f"http://{server_address}/prompt", json=p, headers=headers)
    response.raise_for_status()  # 如果返回错误状态码，则抛出异常
    return response.json()

def get_image(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    response = requests.get(f"http://{server_address}/view", params=data)
    response.raise_for_status()
    return response.content

def get_history(prompt_id):
    response = requests.get(f"http://{server_address}/history/{prompt_id}")
    response.raise_for_status()
    return response.json()

def get_images(ws, prompt):
    prompt_id = queue_prompt(prompt)['prompt_id']
    output_images = {}
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break  # Execution is done
        else:
            continue  # previews are binary data

    history = get_history(prompt_id)[prompt_id]
    for o in history['outputs']:
        for node_id in history['outputs']:
            node_output = history['outputs'][node_id]
            if 'images' in node_output:
                images_output = []
                for image in node_output['images']:
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)
            output_images[node_id] = images_output

    return output_images

if len(sys.argv) != 17:
    print("Usage: python api1.py <flower_image_path> <vanpic_image_path> <subjectpos> <subjectneg> <backgroundpos> <backgroundneg> <subjectseed_1> <subjectseed_2> <subjectseed_3> <subjectseed_4> <backgroundseed_1> <backgroundseed_2> <backgroundseed_3> <backgroundseed_4> <output_directory> <json_file_path>")
    sys.exit(1)

flower_image_path = sys.argv[1]
vanpic_image_path = sys.argv[2]
subjectpos = sys.argv[3]
subjectneg = sys.argv[4]
backgroundpos = sys.argv[5]
backgroundneg = sys.argv[6]
subjectseed_1 = int(sys.argv[7])
subjectseed_2 = int(sys.argv[8])
subjectseed_3 = int(sys.argv[9])
subjectseed_4 = int(sys.argv[10])
backgroundseed_1 = int(sys.argv[11])
backgroundseed_2 = int(sys.argv[12])
backgroundseed_3 = int(sys.argv[13])
backgroundseed_4 = int(sys.argv[14])
output_directory = sys.argv[15]
json_file_path = sys.argv[16]

with open(json_file_path, "r", encoding="utf-8") as f:
    workflow_jsondata = f.read()

prompt = json.loads(workflow_jsondata)

prompt["317"]["inputs"]["image"] = flower_image_path
prompt["227"]["inputs"]["image"] = vanpic_image_path

prompt["225"]["inputs"]["positive"] = subjectpos
prompt["225"]["inputs"]["negative"] = subjectneg
prompt["243"]["inputs"]["positive"] = subjectpos
prompt["243"]["inputs"]["negative"] = subjectneg
prompt["257"]["inputs"]["positive"] = subjectpos
prompt["257"]["inputs"]["negative"] = subjectneg
prompt["271"]["inputs"]["positive"] = subjectpos
prompt["271"]["inputs"]["negative"] = subjectneg

prompt["238"]["inputs"]["positive"] = backgroundpos
prompt["238"]["inputs"]["negative"] = backgroundneg
prompt["252"]["inputs"]["positive"] = backgroundpos
prompt["252"]["inputs"]["negative"] = backgroundneg
prompt["266"]["inputs"]["positive"] = backgroundpos
prompt["266"]["inputs"]["negative"] = backgroundneg
prompt["280"]["inputs"]["positive"] = backgroundpos
prompt["280"]["inputs"]["negative"] = backgroundneg

prompt["226"]["inputs"]["seed"] = subjectseed_1
prompt["244"]["inputs"]["seed"] = subjectseed_2
prompt["258"]["inputs"]["seed"] = subjectseed_3
prompt["272"]["inputs"]["seed"] = subjectseed_4

prompt["239"]["inputs"]["seed"] = backgroundseed_1
prompt["253"]["inputs"]["seed"] = backgroundseed_2
prompt["267"]["inputs"]["seed"] = backgroundseed_3
prompt["281"]["inputs"]["seed"] = backgroundseed_4

ws = websocket.WebSocket()
ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
images = get_images(ws, prompt)

current_time = datetime.now().strftime("%Y%m%d%H%M%S")

# Commented out code to display the output images:
save_node_ids = ["242", "256", "270", "284"]
for node_id in save_node_ids:
    if node_id in images:
        for image_data in images[node_id]:
            from PIL import Image
            import io
            image = Image.open(io.BytesIO(image_data))
            image.show()
            node_title = prompt[node_id]["_meta"]["title"]
            image_path = os.path.join(output_directory, f"{node_title}_{current_time}.png")
            print(f"Saving image at: {image_path}")
            image.save(image_path)